import numpy as np
import matplotlib.patches as patches

class Point:
    """表示二维空间中的一个点"""
    def __init__(self, x, y):
        self.x = x
        self.y = y

class Rectangle:
    """表示二维空间中的一个矩形区域"""

    def __init__(self, x, y, width, height):
        self.x = x  # 左上角x坐标
        self.y = y  # 左上角y坐标
        self.width = width
        self.height = height

    def contains(self, point):
        """检查点是否在矩形内"""
        return (
            self.x <= point.x < self.x + self.width and
            self.y <= point.y < self.y + self.height
        )

    def intersects(self, other):
        """检查两个矩形是否相交"""
        return not (
                other.x > self.x + self.width or
                other.x + other.width < self.x or
                other.y > self.y + self.height or
                other.y + other.height < self.y
        )

class Quadtree:
    """四叉树实现（针对二值网格地图）"""

    def __init__(self, boundary, grid, max_depth=8, depth=0):
        self.boundary = boundary  # 当前节点的边界
        self.max_depth = max_depth  # 最大递归深度
        self.depth = depth  # 当前深度
        self.divided = False  # 是否已分裂
        self.children = []  # 四个子节点
        self.is_obstacle = None  # 当前区域是否全为障碍物
        self.is_buffer_neighbors = False  # build_neighbor_cache()会设置它为True，它决定了get_neighbors的行为。
        self.neighbors = []  # 存储邻居节点
        self.test = {"mask": False, "color": ""} # 用于测试：被标记的网格，需要用特殊颜色绘制。

        # 检查当前区域是否同质（全空白或全障碍）
        x_start, y_start = int(boundary.x), int(boundary.y)
        x_end, y_end = int(boundary.x + boundary.width), int(boundary.y + boundary.height)
        region = grid[y_start:y_end, x_start:x_end]

        if np.all(region == 1):  # 全空白
            self.is_obstacle = False
        elif np.all(region == 0):  # 全障碍
            self.is_obstacle = True
        else:  # 非同质，需要分裂
            if self.depth < self.max_depth:
                self.subdivide(grid)

    def subdivide(self, grid):
        """将当前节点分裂为四个子节点"""
        x, y = self.boundary.x, self.boundary.y
        w, h = self.boundary.width / 2, self.boundary.height / 2

        # 四个象限的边界
        nw = Rectangle(x, y, w, h)
        ne = Rectangle(x + w, y, w, h)
        sw = Rectangle(x, y + h, w, h)
        se = Rectangle(x + w, y + h, w, h)

        # 创建四个子节点
        self.children = [
            Quadtree(nw, grid, self.max_depth, self.depth + 1),
            Quadtree(ne, grid, self.max_depth, self.depth + 1),
            Quadtree(sw, grid, self.max_depth, self.depth + 1),
            Quadtree(se, grid, self.max_depth, self.depth + 1)
        ]
        self.divided = True

    def is_passable(self):
        """检查节点是否可通行（非障碍物且为叶节点）"""
        return not self.divided and self.is_obstacle is not None and not self.is_obstacle

    def get_neighbors(self, quadtree_root):
        """获取当前节点相邻的所有可通行叶节点"""
        if not self.is_buffer_neighbors:
            self.neighbors.clear()
            self._find_adjacent_leaves(quadtree_root, self.boundary)
        return self.neighbors

    def _find_adjacent_leaves(self, node, target_boundary):
        """递归查找与目标边界相邻的叶节点"""
        if not node.boundary.intersects(target_boundary):
            return  # 无交集则跳过

        if node.is_passable() and node != self:
            self.neighbors.append(node)  # 找到可通行的相邻叶节点
        elif node.divided:
            for child in node.children:
                self._find_adjacent_leaves(child, target_boundary)

    def build_neighbor_cache(self, root):
        """递归构建邻居缓存"""
        self.is_buffer_neighbors = True
        if self.is_passable():
            self._find_adjacent_leaves(root, self.boundary)
        elif self.divided:
            for child in self.children:
                child.build_neighbor_cache(root)

    def draw(self, ax):
        """绘制四叉树结构"""
        if self.is_obstacle is not None:  # 叶节点
            color = 'black' if self.is_obstacle else 'white'
            if self.test["mask"] and self.test["color"] != "":  # 测试用
                color = self.test["color"]
            rect = patches.Rectangle(
                (self.boundary.x, self.boundary.y),
                self.boundary.width,
                self.boundary.height,
                facecolor=color,
                edgecolor='gray',
                linewidth=0.5
            )
            ax.add_patch(rect)
        elif self.divided:  # 非叶节点
            for child in self.children:
                child.draw(ax)